import cv2
import numpy as np
import math
from math import *

from skimage.measure import label


def dilate_demo(image, size):
    # 膨胀操作
    ret, binary = cv2.threshold(image, 0, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (size, size))  # 定义结构元素的形状和大小
    erd = cv2.erode(binary, kernel)
    return erd


def erode_demo(image, size):
    # 腐蚀操作
    ret, binary = cv2.threshold(image, 0, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (size, size))  # 定义结构元素的形状和大小
    dlt = cv2.dilate(binary, kernel)
    return dlt


def rotate_img(image, degree):
    height, width = image.shape[:2]
    image = image.astype("uint8")
    height_new = int(width * fabs(sin(radians(degree))) + height * fabs(cos(radians(degree))))
    width_new = int(height * fabs(sin(radians(degree))) + width * fabs(cos(radians(degree))))

    mat_rotation = cv2.getRotationMatrix2D((width / 2, height / 2), degree, 1)
    mat_rotation[0, 2] += (width_new - width) / 2
    mat_rotation[1, 2] += (height_new - height) / 2
    return cv2.warpAffine(image, mat_rotation, (width_new, height_new), borderValue=(255, 255, 255))


def get_rect(label_image):
    res = [0, 0, 0, 0]
    height, width = label_image.shape
    find_one = False
    area = 0
    for i in range(height):
        for j in range(width):
            if label_image[i][j] == 0:
                area += 1
                if not find_one:
                    find_one = True
                    res = [j, j, i, i]
                else:
                    if j < res[0]:
                        label_image[i][j] = 120
                        res[0] = j
                    if j > res[1]:
                        res[1] = j
                    if i < res[2]:
                        res[2] = i
                    if i > res[3]:
                        res[3] = i
    return res, area / ((res[1] - res[0]) * (res[3] - res[2]))


def get_horizontal_shadow(img_bi):  # 水平投影+分割
    # 1.水平投影
    h, w = img_bi.shape
    shadow_h = img_bi.copy()  # shadow_h画图用(切记！copy后面还有个())

    a = [0 for z in range(0, h)]  # 初始化一个长度为h的数组，用于记录每一行的黑点个数

    for j in range(0, h):  # 遍历一行
        for i in range(0, w):  # 遍历一列
            if shadow_h[j, i] < 255:  # 发现黑色
                a[j] += 1  # a数组这一行的值+1
                shadow_h[j, i] = 255  # 记录好了就变为白色

    for j in range(0, h):  # 遍历一行 画黑条,长度为a[j]
        for i in range(0, a[j]):
            shadow_h[j, i] = 0
    # cv2.imshow('shadow_'+str(h), shadow_h)

    # 计算a的方差

    start_index = 0
    end_index = len(a) - 1
    for i in range(len(a)):
        if a[i] != 0:
            start_index = i
            break

    for i in range(0, len(a)):
        if a[-(i + 1)] != 0:
            end_index = -(i + 1)
            break
    a = a[start_index:end_index]
    arr = np.array(a)
    return shadow_h, np.var(arr)


# 度数转换
def DegreeTrans(theta):
    res = theta / np.pi * 180
    return res


# 通过霍夫变换计算角度
def CalcDegree(srcImage):
    dstImage = cv2.Canny(srcImage, 50, 200, 3)
    lineimage = srcImage.copy()
    height = srcImage.shape[0]

    # 通过霍夫变换检测竖向直线
    lines = []
    for i in range(50, 200, 1):
        lines = cv2.HoughLines(dstImage, 1, np.pi / 180, i)
        # print(len(lines))
        if len(lines) == 1:
            break
    # 由于图像不同，阈值不好设定，因为阈值设定过高导致无法检测直线，阈值过低直线太多，速度很慢
    sum = 0
    # 依次画出每条线段
    for i in range(len(lines)):
        for rho, theta in lines[i]:
            a = np.cos(theta)
            b = np.sin(theta)
            x0 = a * rho
            y0 = b * rho
            x1 = int(round(x0 + 1000 * (-b)))
            y1 = int(round(y0 + 1000 * a))
            x2 = int(round(x0 - 1000 * (-b)))
            y2 = int(round(y0 - 1000 * a))
            sum += theta
            cv2.line(lineimage, (x1, y1), (x2, y2), (0, 0, 255), 1, cv2.LINE_AA)

    cv2.imshow("I"+str(height), lineimage)
    average = sum / len(lines)
    angle = DegreeTrans(average)
    return angle


def Warp(image, angle):
    angle = -angle
    if abs(90-abs(angle)) < abs(180-abs(angle)) and abs(90-abs(angle)) < abs(angle):
        angle = angle+90
    a = math.tan(angle * math.pi / 180.0)
    height = image.shape[0]
    width = int(image.shape[1] + height * a)

    iWarp = np.zeros((height, width), dtype=np.uint8)
    iWarp = cv2.cvtColor(iWarp, cv2.COLOR_GRAY2BGR)
    for i in range(image.shape[0]):
        for j in range(image.shape[1]):
            x = int(j + i * a)
            if x < width:
                tmp = image[i, j]
                iWarp[i, x] = tmp
    return iWarp


def pickBlue(img, scale):
    res = []
    img_origin = img.copy()
    h, w, _ = img.shape
    cv2.imshow('origin', img)

    print("将蓝色部分提取出来")
    border_img = cv2.Canny(img, 50, 150)
    bh, bw = border_img.shape[:2]
    for j in range(bh):
        for k in range(bw):
            border_img[j][k] = 255 - border_img[j][k]
    border_img = dilate_demo(border_img, 2)

    for i in range(0, h):
        for j in range(0, w):
            B = float(img[i][j][0])
            G = float(img[i][j][1])
            R = float(img[i][j][2])
            if (B / max(G + R + B, 1)) * 255 > 86 or R + G + B < 240:
                area_sum = 0
                canny_sum = 0
                for x in range(i-2, i+2, 1):
                    for y in range(j-2, j+2):
                        if x < h and y < w:
                            area_sum += 1
                            if border_img[x, y] == 0:
                                canny_sum += 1
                if canny_sum/area_sum < 0.7:
                    img[i][j] = [255, 255, 255]
            else:
                img[i][j] = [255, 255, 255]

    cv2.imshow('pickBlue', img)
    cv2.imwrite("cutLines/" + "pickblue.jpg", img)
    cv2.imshow("borderImage", border_img)

    # 进行二值化，并膨胀腐蚀

    print("进行二值化，并膨胀腐蚀")
    img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    img = erode_demo(img, 2)
    cv2.imshow('after_erode', img)
    img = dilate_demo(img, scale + 2)
    cv2.imshow('after_dilate', img)
    cv2.imwrite('cutLines/after_dilate.jpg', img)

    print("计算不同的连通域，将面积最大的几块取出来，认为是要取出的几块")
    labeled_img, num = label(img, connectivity=1, background=1, return_num=True)

    label_sum = [0] * num
    label_num = [0] * num
    for i in range(num):
        label_num[i] = i + 1

    for i in range(0, h):
        for j in range(0, w):
            label_sum[labeled_img[i][j] - 1] += 1

    for i in range(num):
        for j in range(i + 1, num):
            if label_sum[i] < label_sum[j]:
                tmp = label_sum[i]
                label_sum[i] = label_sum[j]
                label_sum[j] = tmp
                tmp = label_num[i]
                label_num[i] = label_num[j]
                label_num[j] = tmp

    break_point = 0

    for i in range(1, num - 1):
        if label_sum[i + 1] < label_sum[1] / 3:
            break_point = i
            break
    chosen_label = label_num[1:break_point + 1]
    for i in range(0, h):
        for j in range(0, w):
            if labeled_img[i][j] in chosen_label:
                img[i][j] = 0
            else:
                img[i][j] = 255
    cv2.imshow('pick_biggest_part', img)
    print("共计选出：", len(chosen_label), "个区域")

    print("将最大的几块连通区域单独处理:")
    for i in chosen_label:
        single_labeled_img = labeled_img.copy()
        for j in range(h):
            for k in range(w):
                if labeled_img[j][k] == i:
                    single_labeled_img[j][k] = 0
                else:
                    single_labeled_img[j][k] = 255

        max_area_percent = 0
        max_rect = [0, 0, 0, 0]
        max_rotate = 0
        print("\n为", i, "号区寻找最佳旋转角度:")
        for j in range(-10, 10, 1):
            r_img = rotate_img(single_labeled_img, j)
            rect, area_percent = get_rect(r_img)
            if area_percent > max_area_percent:
                max_area_percent = area_percent
                max_rect = rect
                max_rotate = j
            print("--", end='')
        print("\n图像切割选择角度:", max_rotate)
        r_img_origin = rotate_img(img_origin, max_rotate)
        new_single_labeled_img = np.zeros((max_rect[3] - max_rect[2] + 1, max_rect[1] - max_rect[0] + 1),
                                          dtype=np.uint8)
        if new_single_labeled_img.shape[0] < new_single_labeled_img.shape[1]:
            print("此部分因长宽比不合适被舍去")
            continue
        new_single_labeled_img = cv2.cvtColor(new_single_labeled_img, cv2.COLOR_GRAY2BGR)
        for j in range(max_rect[2], max_rect[3] + 1):
            for k in range(max_rect[0], max_rect[1] + 1):
                new_single_labeled_img[j - max_rect[2]][k - max_rect[0]] = r_img_origin[j][k]

        cv2.imshow('single_l_' + str(i), new_single_labeled_img)
        # 以下对切出来的图像进行边缘处理，然后准备进行切斜检测
        max_var = 0
        hori_degree = 0

        # 寻找单行平行线
        border_img = cv2.Canny(cv2.GaussianBlur(new_single_labeled_img, (3, 3), 0), 50, 150)
        h_single, w_single = border_img.shape[:2]
        for j in range(h_single):
            for k in range(w_single):
                border_img[j][k] = 255 - border_img[j][k]
        for j in range(-10, 10, 1):
            r_img = rotate_img(border_img, j)
            canny, var_r = get_horizontal_shadow(r_img)
            if var_r > max_var:
                max_var = var_r
                hori_degree = j
        print("文本平行线选择角度：", hori_degree)

        border_img = rotate_img(border_img, hori_degree)
        hough_img = dilate_demo(border_img, 12)
        vect_degree = CalcDegree(hough_img)
        print("文本垂直线选择角度：", vect_degree)

        new_single_labeled_img = rotate_img(new_single_labeled_img, hori_degree)
        warp_img = Warp(new_single_labeled_img, 180 - vect_degree)
        res.append(warp_img)
        cv2.imshow("w_l_" + str(i), warp_img)
    return res
